use clap::Parser;
use serde_json::{Map, Value};
use std::{collections::HashMap, fs::File, io::Write, path::PathBuf};
use tokenizers::{Encoding, Tokenizer};

pub mod config;
pub mod model;
pub mod params;
pub mod tensor;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Path to the Model-dir
    #[arg(short, long)]
    model_path: Option<PathBuf>,

    /// Path to the output file
    #[arg(short, long)]
    output_path: Option<PathBuf>,

    /// Input text to tokenize
    #[arg(long, default_value_t = String::from("Once upon a time"))]
    input: String,
}
fn main() {
    let args = Args::parse();
    let model_path = args.model_path.unwrap_or_else(get_default_model_path);
    let output_dir = args.output_path.unwrap_or_else(get_default_output_path);

    let tokenizer_file_path = model_path.join("tokenizer.json");
    let tokens = tokenize(
        &tokenizer_file_path,
        args.input,
        Some(&output_dir.join("vocab.json")),
    );
    println!("Tokens: {:?}", tokens);

    save_model(&model_path, &output_dir.join("example.json"));
}

fn get_default_model_path() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .join("models")
        .join("story")
}

fn get_default_output_path() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .join("results")
}

fn tokenize(file_path: &PathBuf, input: String, output_path: Option<&PathBuf>) -> Encoding {
    let tokenizer = Tokenizer::from_file(file_path).unwrap();
    match output_path {
        None => {}
        Some(output_path) => {
            let vocab = tokenizer.get_vocab(true);
            save_vocab(&vocab, output_path);
        }
    }
    tokenizer.encode(input, true).unwrap()
}

fn save_model(input_path: &PathBuf, output_path: &PathBuf) {
    let llama = model::Llama::<f32>::from_safetensors(input_path);

    if output_path.exists() {
        // Handle the case where the output path already exists
        println!("Output path already exists...  Skipped saving model.");
    } else {
        let mut file = File::create(output_path).unwrap();
        let _ = file.write_all(llama.to_json().unwrap().to_string().as_bytes());
        println!("Model saved to {:?}", output_path);
    }
}

fn save_vocab(vocab: &HashMap<String, u32>, output_path: &PathBuf) {
    let mut vocab_json = Map::new();
    for (key, value) in vocab {
        vocab_json.insert(key.to_string(), Value::from(*value));
    }

    let vocab_str = serde_json::to_string_pretty(&vocab_json).unwrap();
    if output_path.exists() {
        println!("Output path already exists... Skipped saving vocab.");
        return;
    }

    let mut file = File::create(output_path).unwrap();
    file.write_all(vocab_str.as_bytes()).unwrap();
    println!("Vocab saved to {:?}", output_path);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_tokenize() {
        let file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .join("models")
            .join("story")
            .join("tokenizer.json");

        let input = String::from("Hello, World!");
        let tokens = tokenize(&file_path, input, None);
        assert!(
            tokens.get_ids().len() > 0,
            "The number of token IDs should be greater than 0."
        );
    }

    #[test]
    fn test_tokenize_with_example_file() {
        let input = String::from("Hello, World!");
        use_example_dir(|file_path| {
            let tokens = tokenize(&file_path.join("tokenizer.json"), input.clone(), None);
            assert!(
                tokens.get_ids().len() > 0,
                "The number of token IDs should be greater than 0."
            );
            println!("Tokens: {:?}", tokens);
        });
    }

    #[test]
    fn test_use_example_dir_print() {
        use_example_dir(|path| {
            println!("path: {:#?}", path);
            assert!(path.exists(), "The example file path should exist.");
            assert!(
                path.is_dir(),
                "The example file path should point to a file."
            );
        });
    }
}

pub fn use_example_dir<F>(f: F)
where
    F: FnOnce(&PathBuf),
{
    let file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .join("models")
        .join("story");
    f(&file_path);
}
